import sys
sys.path.append("..")
from util import multiprocess,load_json_file
import numpy as np
from tqdm import tqdm
import torch
from transformers import BertModel
from tqdm import tqdm
import numpy as np
import time

device=torch.device("cuda:0")
bert_model=BertModel.from_pretrained("/home/yujia_zhou/long_doc/anchors-master/output/bert/",output_attentions=True)
bert_model=bert_model.to(device)

def get_cos_similar(v1, v2):
    num = float(np.dot(v1, v2))  # 向量点乘
    denom = np.linalg.norm(v1) * np.linalg.norm(v2)  # 求模长的乘积
    return 0.5 + 0.5 * (num / denom) if denom != 0 else 0 # 将余弦相似度转换到[0,1]

# 计算向量a,b的余弦相似度
def get_att_dis(a, b):
    with torch.no_grad():
        ans = []
        a,b=torch.tensor(np.array(a)).to(device),torch.tensor(np.array(b)).to(device)
        a,b=bert_model.embeddings.word_embeddings(a),bert_model.embeddings.word_embeddings(b)     
        a=torch.sum(a,dim=0)/a.shape[0]

        t_a=a.cpu().numpy()
        t_b=b.cpu().numpy()

        for i in range(len(b)):
            attention_score=get_cos_similar(t_a, t_b[i])
            ans.append(attention_score)
        return ans

# 处理包含正负例的数据
def get_cos_weight_posneg(dic):
    qry=dic['qry']['query']
    psgs={}
    psgs['pos']=[dic['pos'][0]['passage']]
    psgs['neg']=[]
    for i in range(7):
        psgs['neg'].append(dic['neg'][i]['passage'])

    ans={}
    for type_ in psgs.keys():
        print("---------------正在处理：",type_,"----------------------")
        for psg in psgs[type_]:
            with torch.no_grad():
                w=get_att_dis(qry,psg)
                if type_ not in ans.keys():
                    ans[type_]=[]              
                w=np.array(w)
                cur=list(w.shape[0]*w/np.sum(w))
                ans[type_].append(cur)
    return ans 


# 处理只包含正例的数据
def get_cos_weight_single(dic):
    qry=dic['qry']['query']
    psgs={}
    psgs['pos']=[dic['pos'][0]['passage']]
    psgs['neg']=[]
    for i in range(7):
        psgs['neg'].append(dic['neg'][i]['passage'])

    ans={}
    for type_ in psgs.keys():
        print("---------------正在处理：",type_,"----------------------")
        for psg in psgs[type_]:
            with torch.no_grad():
                w=get_att_dis(qry,psg)
                if type_ not in ans.keys():
                    ans[type_]=[]              
                w=np.array(w)
                cur=list(w.shape[0]*w/np.sum(w))
                ans[type_].append(cur)
    return ans 


if __name__=="__main__":
    torch.multiprocessing.set_start_method('spawn')
    device=torch.device("cuda:0")
    # device=torch.device("cpu")
    for t in range(100):
        s="%03d"%t
        data,len_li=load_json_file(root_path="/home/huaying_yuan/long/data/lce_data/"+s+".group.json")
        w_dynamic_list=multiprocess(get_cos_weight_posneg,data,pool_size=8)

        with open("/home/huaying_yuan/long/weights/lce_data/cos/"+s+".txt","w") as f:
            for i in tqdm(w_dynamic_list):
                print(i,file=f)

